package com.sz.admin.ai.util;

import lombok.SneakyThrows;
import lombok.extern.log4j.Log4j2;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.JsonReader;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.ParagraphPdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.core.io.Resource;
import org.springframework.core.io.UrlResource;
import org.springframework.stereotype.Component;

import java.net.MalformedURLException;
import java.util.List;

/**
 * @描述: 向量工具类
 * @Author: TangYuan
 * @Date: 2025/3/3 10:35
 */
@Log4j2
@Component
public class EmbeddingUtil {
    
    /**
     * 根据url读取文档并拆分
     * @param fileUrl 文件url
     * @return {@link List<Document>}
     */
    @SneakyThrows
    public static List<Document> readAndSplit(String fileUrl)  {
        log.info("----------- 开始读取文档");
        String[] split = fileUrl.split("\\.");
        String suffix = split[split.length - 1];
        
        List<Document> documentList = switch (suffix) {
            // 从JSON中读取器
            case "json" -> jsonReader(fileUrl);
            // 从文本中读取器
            case "txt" -> textReader(fileUrl);
            // 从PDF中读取器
            case "pdf" -> pdfParagraphReader(fileUrl);
            // 通用MD读取器
            case "md" -> markdownReader(fileUrl);
            // 通用文档读取器
            default -> tikaReader(fileUrl);
        };
        documentList.forEach(document -> {
            System.out.println();
            System.out.println("---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------");
            System.out.println(document.getText());
            System.out.println("---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------");
            System.out.println();
        });
        
        return documentList;
    }
    
    /**
     * 从JSON中读取器
     * @param fileUrl fileUrl
     * @return {@link List<Document>}
     */
    private static List<Document> jsonReader(String fileUrl) throws MalformedURLException {
        log.info("----------- jsonReader");
        Resource resource = new UrlResource(fileUrl);
        JsonReader jsonReader = new JsonReader(resource);
        return jsonReader.read();
    }
    
    /**
     * 从文档中读取器
     * @param fileUrl fileUrl
     * @return {@link List<Document>}
     */
    private static List<Document> textReader(String fileUrl) {
        log.info("----------- textReader");
        TextReader textReader = new TextReader(fileUrl);
        return textReader.read();
    }
    
    /**
     * 从PDF中读取器
     * @param fileUrl fileUrl
     * @return {@link List<Document>}
     */
    private static List<Document> pdfReader(String fileUrl) {
        log.info("----------- pdfReader");
        PdfDocumentReaderConfig build = PdfDocumentReaderConfig
                .builder()
                .withPagesPerDocument(1)
                .build();
        PagePdfDocumentReader pagePdfDocumentReader = new PagePdfDocumentReader(fileUrl, build);
        return pagePdfDocumentReader.read();
    }
    
    /**
     * 从PDFParagraph中读取器
     * @param fileUrl fileUrl
     * @return {@link List<Document>}
     */
    private static List<Document> pdfParagraphReader(String fileUrl) {
        log.info("----------- pdfParagraphReader");
        List<Document> document = null;
        try {
            PdfDocumentReaderConfig build = PdfDocumentReaderConfig
                    .builder()
                    .withPagesPerDocument(1)
                    .build();
            ParagraphPdfDocumentReader paragraphPdfDocumentReader = new ParagraphPdfDocumentReader(fileUrl, build);
            document = paragraphPdfDocumentReader.read();
        }catch (IllegalArgumentException e) {
            log.error("没有TOC目录，转为使用按页读取 -----> 异常：{}", e.getMessage());
            document = pdfReader(fileUrl);
        }
        
        return document;
    }
    
    /**
     * 通用文档读取器
     * @param fileUrl fileUrl
     * @return {@link List<Document>}
     */
    private static List<Document> markdownReader(String fileUrl) {
        log.info("----------- markdownReader");
        MarkdownDocumentReaderConfig markdownDocumentReaderConfig = MarkdownDocumentReaderConfig
                .builder()
                .withHorizontalRuleCreateDocument(true)
                .withIncludeBlockquote(true)
                .build();
        MarkdownDocumentReader markdownDocumentReader = new MarkdownDocumentReader(fileUrl, markdownDocumentReaderConfig);
        return markdownDocumentReader.read();
    }
    
    /**
     * 通用文档读取器
     * @param fileUrl fileUrl
     * @return {@link List<Document>}
     */
    private static List<Document> tikaReader(String fileUrl) {
        log.info("----------- tikaReader");
        TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(fileUrl);
        List<Document> document = tikaDocumentReader.read();
        TokenTextSplitter tokenTextSplitter = TokenTextSplitter.builder()
                .withChunkSize(800)
                .withMinChunkSizeChars(350)
                .withMinChunkLengthToEmbed(5)
                .withMaxNumChunks(1000)
                .withKeepSeparator(true)
                .build();
        return tokenTextSplitter.apply(document);
    }
}
